from openai import OpenAI
import os

def call_openai_style_api(api_base, api_key, model, prompt, temperature=0.3, max_tokens=4000, top_p=0.6):
    client = OpenAI(
        api_key=api_key,
        base_url=api_base
    )
    response = client.chat.completions.create(
        model=model,
        messages=[{
            "role": "user",
            "content": prompt,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "top_p": top_p,
        }]
    )

    print(response.choices[0].message.content)
    return response.choices[0].message.content

def call_doubao_api(prompt, temperature=0.3, max_tokens=4000, top_p=0.6):
    api_key=os.environ.get("ARK_API_KEY"),
    api_base="https://ark.cn-beijing.volces.com/api/v3",
    model=os.environ.get("ARK_ENDPOINT_ID"),
    response = call_openai_style_api(api_base, api_key, model, prompt, temperature, max_tokens, top_p)

def call_lingyiwanwu_api(prompt, temperature=0.3, max_tokens=4000, top_p=0.6):
    api_base= "https://api.lingyiwanwu.com/v1"
    api_key = os.environ.get("LINGYIWANWU_API_KEY")

    response = call_openai_style_api(api_base, api_key, "yi-lightning", prompt, temperature, max_tokens, top_p)
    return response
    # client = OpenAI(
    #     api_key=API_KEY,
    #     base_url=API_BASE
    # )
    # completion = client.chat.completions.create(
    #     model="yi-lightning",
    #     messages=[{
    #         "role": "user",
    #         "content": prompt,
    #         "temperature": temperature,
    #         "max_tokens": max_tokens,
    #         "top_p": top_p,
    #     }]
    # )
    #
    # return completion.choices[0].message.content



def call_zhipu_api(prompt, temperature=0.8, max_tokens=4095, top_p=0.6):
    from zhipuai import ZhipuAI
    api_key = os.environ.get("ZHIPU_API_KEY")
    # print(api_key)
    client = ZhipuAI(api_key=api_key)
    response = client.chat.completions.create(
        model="glm-4-flash",
        messages=[
            {"role": "user", "content": prompt},
        ],
        temperature=temperature,
        max_tokens=max_tokens,
        top_p=top_p
    )
    print(response.choices[0].message.content
          )
    return response.choices[0].message.content

def call_deepseek_api(prompt, temperature=0.8, max_tokens=8000, top_p=0.6):
    api_key = os.environ.get("DEEPSEEK_API_KEY")
    api_base = "https://api.deepseek.com"

    response = call_openai_style_api(api_base, api_key, "deepseek-chat", prompt, temperature, max_tokens, top_p)
    return response

